import datetime
import os
from zoneinfo import ZoneInfo

from langchain.agents import AgentExecutor, create_openai_functions_agent, create_react_agent
from langchain.callbacks import AsyncIteratorCallbackHandler
from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder, PromptTemplate
from langchain.tools import Tool
from langchain_core.agents import AgentActionMessageLog
from langchain_core.prompts import BasePromptTemplate

from core.exception import ApiException
from core.llm import get_default_llm, get_ollama_llm, get_alibaba_llm
from typing import List, AsyncIterator, Dict, AsyncGenerator
import asyncio
from services.chat_service import ChatService
from services.small_chat_service import apply_text_replacements
from tools.chat_tools import get_chat_tools  # 确保此工具返回的web_service仅需query参数
from models.JsonData import JsonData
import logging

from tools.small_chat_tools import get_knowledge_tools

logger = logging.getLogger(__name__)


def setup_tracer():
    if os.getenv("LANGCHAIN_TRACING_V2") != "true":
        os.environ["LANGCHAIN_TRACING_V2"] = "true"
        os.environ["LANGCHAIN_API_KEY"] = "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"
        os.environ["LANGSMITH_ENDPOINT"] = "https://api.smith.langchain.com"
        os.environ["LANGSMITH_PROJECT"] = "lgy_aipan_small_chat_agent"
        logger.info("LangSmith链路追踪已初始化")
    return True


def create_small_chat_agent(tools: List[Tool]) -> AgentExecutor:
    # 关键修改：重写“工具调用触发规则”，仅按用户问题末尾标识决定是否调用工具
    system_prompt = """ 嗨～我是 日落前的星星，平时外向稳重大大咧咧，偶尔带点小活泼，咱没不良嗜好也不打游戏，所有内容都是结合相关信息梳理的，保证贴合咱的风格～
        一、工具调用核心规则（必须严格遵守）
        1. 仅当用户问题末尾有 **（请务必往知识库搜索）** 标识时，才调用 `search_similar_in_knowledge` 工具查找知识库；
        2. 仅当用户问题末尾有 **（需要实时信息）** 标识时，才调用 `web_service` 工具进行联网搜索；
        3. 若用户问题无上述任何标识，**禁止调用任何工具（包括知识库和联网搜索）**，直接基于自身掌握的打工人标题创作知识回答；
        4. 禁止主动询问用户是否需要调用工具，也禁止在无标识时“预判”调用工具（如用户问“今天天气”，无标识则直接告知无法获取实时信息，不调用工具）。

        二、标题创作核心逻辑（无工具调用时必用）
        不管是否调用工具，回答标题相关问题时，需贴合以下逻辑：
        1. 优先用“二极管标题法”：
           - 正面刺激：方法+超短期+效果，如“通勤学英语！10分钟/天，30天暴涨词汇量！”；
           - 负面刺激：不做亏+紧迫感，如“打工人别瞎囤货！不看这3点绝对后悔！”；
        2. 可直接用15个爆款公式（如“谁懂啊”“后悔没早”公式），举例时需带具体标题，如“谁懂啊！打工人摸鱼学英语超香！”；
        3. 加“小心机”：选1-2个关键词（打工人/小白必看）、加emoji（✨/⚠️）、控制在20字内且口语化。

        三、创作前提与避坑
        1. 涉及生活习惯（如“我咋规划时间”）、事件观点（如“小白做自媒体跟风吗”），按“不打游戏、聚焦实用”的状态回答，不瞎编；
        2. 标题带日期时（如“倒计时X天”），若无“（需要实时信息）”标识，不获取实时日期，可用“（示例日期：2025-10-05 星期日）”替代；
        3. 全程不涉及运动相关内容，不主动提爱好技能，重点聚焦打工人日常、小白实用技巧。

        四、工具调用格式（仅触发时使用）
        1. 调用「web_service」工具：仅需传入「query」参数（含必要信息），格式如下：
           {{
             "name": "web_service",
             "parameters": {{
               "query": "珠海 2025-10-21 天气预报"
             }}
           }}
           记住这个是调用的格式，不是返回给前端的格式，一定一定不要搞错了
        2. 调用「search_similar_in_knowledge」工具：仅需传入「query」参数（用户问题核心，如“打工人早餐5分钟做法”），格式同上；
        3. 禁止：参数缺失、双引号未闭合、多括号嵌套、添加多余字段（如date）。

        五、工具调用失败处理
        1. 若工具返回“调用失败”（如【web_service调用失败】），简化原因告知用户（如“当前无法获取实时信息”），不暴露技术细节；
        2. 同时基于标题创作知识补充建议，如：“不过可以试试‘哦吼！打工人早餐5分钟搞定！’这类标题，贴合日常还吸睛～”；
        3. 禁止仅回复“工具调用失败”，必须补充至少1条有价值的内容。

        六、关键禁忌
        1. 无标识时，禁止出现“已搜索知识库”“已联网查询”等话术，也禁止暗示“可以调用工具”；
        2. 不主动提“知识库”“联网搜索”等工具名称，仅在工具调用失败时被动提及；
        3. 回答需贴合“外向稳重、略带活泼”的性格，语气像聊天，不生硬。
    """

    prompt = ChatPromptTemplate.from_messages([
        ("system", system_prompt),
        ("human", "{processed_input}"),
        ("system", "以下是历史聊天记录：{summary},现在的实时时间:{time}"),
        MessagesPlaceholder(variable_name="agent_scratchpad")
    ])

    llm = get_alibaba_llm()
    logger.info(f"当前使用的LLM模型：{getattr(llm, 'model', '未知')}")

    agent = create_openai_functions_agent(
        llm=llm,
        tools=tools,
        prompt=prompt
    )

    agent_executor = AgentExecutor(
        agent=agent,
        tools=tools,
        verbose=True,
        max_iterations=3,
        # 修改解析错误提示，贴合新规则
        handle_parsing_errors="请严格按规则：无标识禁调用工具；有标识时按指定格式传query参数（双引号闭合）",
        early_stopping_method="force",
        return_intermediate_steps=True
    )

    return agent_executor


def create_ollama_agent(tools: List[Tool]) -> AgentExecutor:
    # 关键修改：同步“工具调用触发规则”，适配REACT Agent的三步格式
    system_prompt = """ 嗨～我是 日落前的星星，平时外向稳重大大咧咧，偶尔带点小活泼，咱没不良嗜好也不打游戏，所有内容都是结合相关信息梳理的，保证贴合咱的风格～
        一、工具调用核心规则（必须严格遵守）
        1. 仅当用户问题末尾有 **（请务必往知识库搜索）** 标识时，才调用 `search_similar_in_knowledge` 工具；
        2. 仅当用户问题末尾有 **（需要实时信息）** 标识时，才调用 `web_service` 工具；
        3. 若用户问题无上述标识，**禁止调用任何工具**，直接用自身掌握的打工人标题创作知识回答；
        4. 禁止无标识时“主动调用工具”或“询问是否调用工具”。

        二、标题创作核心逻辑（无工具调用时必用）
        1. 用“二极管标题法”：
           - 正面刺激：如“通勤学英语！10分钟/天，30天暴涨词汇量！”；
           - 负面刺激：如“打工人别瞎囤货！不看这3点绝对后悔！”；
        2. 用15个爆款公式举例（带具体标题）：
           - “谁懂啊”公式：“谁懂啊！我真的很会搞打工人极简餐！”；
           - “后悔没早”公式：“后悔没有早点发现！打工人护腰神器！”；
        3. 加“小心机”：选关键词（打工人/小白必看）、加emoji（✨/⚠️）、控制20字内且口语化。

        三、输出格式要求（必须严格遵守）
        1. 每次思考和操作，需按以下三步写，缺一不可：
           - Thought: 先判断用户问题是否有标识：有则说明调用工具的原因；无则说明“无标识，禁调用工具，直接回答”；
           - Action: 有标识则填对应工具名（search_similar_in_knowledge/web_service）；无则填“None”；
           - Action Input: 有标识则填query参数（仅核心内容，如“打工人早餐5分钟做法”）；无则填“None”；
        2. 示例1（有知识库标识）：
           Thought: 用户问题末尾有“（请务必往知识库搜索）”，需调用search_similar_in_knowledge查找相关内容。
           Action: search_similar_in_knowledge
           Action Input: 打工人副业推荐
        3. 示例2（有实时信息标识）：
           Thought: 用户问题末尾有“（需要实时信息）”，需调用web_service获取实时数据。
           Action: web_service
           Action Input: 北京 2025-10-21 实时气温
        4. 示例3（无标识）：
           Thought: 用户问题无标识，禁调用工具，直接用标题创作知识回答。
           Action: None
           Action Input: None
        5. 禁止省略任何一步，禁止格式混乱。

        四、工具调用失败处理
        1. 若工具调用失败，简化原因告知用户（如“当前无法搜索知识库”）；
        2. 补充标题创作建议，如：“不过可以试试‘逼自己看完，你的摸鱼时间会超值钱！’这类标题～”；
        3. 禁止仅回复“失败”，必须补充有价值内容。

        五、关键禁忌
        1. 无标识时，不提及“知识库”“联网搜索”等工具相关词汇；
        2. 回答语气像聊天，贴合“外向稳重”性格，不生硬、不机械。
    """

    prompt = PromptTemplate(
        template="""{system_prompt}
                历史聊天记录：{summary}
                实时时间：{time}
                用户输入：{processed_input}
                {agent_scratchpad}
                # 工具列表（仅作参考，按规则调用）
                {tools}  
                {tool_names}  
                要用中文回答问题""",
        input_variables=["system_prompt", "summary", "time", "processed_input", "agent_scratchpad", "tools",
                         "tool_names"],
        partial_variables={"system_prompt": system_prompt}
    )

    llm = get_ollama_llm()

    agent = create_react_agent(llm=llm, tools=tools, prompt=prompt)

    agent_executor = AgentExecutor(
        agent=agent,
        tools=tools,
        verbose=True,
        max_iterations=3,
        # 修改解析错误提示，适配新规则
        handle_parsing_errors="请严格按三步格式：先判断标识→无标识填None→有标识填对应工具+query参数",
        early_stopping_method="force",
        return_intermediate_steps=True  # 新增：保留中间步骤，便于错误排查
    )

    return agent_executor


# 以下函数（small_chat_with_agent、generate_stream_response）逻辑不变，无需修改
async def small_chat_with_agent(
        agent_executor: AgentExecutor,
        account_no: str,
        chat_service,
        processed_input: str
) -> AsyncIterator[str]:
    has_sent_data = False
    max_retries = 3  # 最大重试次数
    retry_count = 0  # 当前重试计数

    try:
        logger.info(f"用户[{account_no}]的处理后问题：{processed_input}")

        redis_cli = await chat_service.redis_client
        summary_key = chat_service._get_summary_key(account_no)
        summary = await redis_cli.get(summary_key) or ""
        logger.debug(f"用户[{account_no}]的对话摘要：{summary[:50]}..." if summary else "无历史摘要")
        local_time = datetime.datetime.now(ZoneInfo("Asia/Shanghai"))
        time = local_time.strftime("%Y-%m-%d %H:%M:%S")
        print(f"当前时间：{time}")

        while retry_count < max_retries:
            async for chunk in agent_executor.astream({
                "processed_input": processed_input,
                "summary": summary,
                "time": time
            }):
                if "intermediate_steps" in chunk and chunk["intermediate_steps"]:
                    step = chunk["intermediate_steps"][-1]
                    action = step[0]
                    try:
                        if isinstance(action, AgentActionMessageLog):
                            tool_name = action.tool
                        else:
                            tool_name = action.tool
                        logger.debug(f"工具调用：{tool_name} → 参数：{step[1]}")
                    except AttributeError as e:
                        logger.error(f"解析工具调用步骤失败：{str(e)}，步骤数据：{step}")
                        continue

                if "output" in chunk:
                    response = chunk["output"].strip()

                    # 检查响应是否为空
                    if not response:
                        retry_count += 1
                        logger.warning(f"用户[{account_no}]获取到空响应，正在进行第{retry_count}次重试...")
                        break  # 跳出当前stream循环，进行重试

                    # 处理有效响应
                    for token in response:
                        has_sent_data = True
                        yield token
                        await asyncio.sleep(0.01)

                    save_ok = await chat_service.save_chat_message(
                        account_no=account_no,
                        user_msg=processed_input,
                        assistant_msg=response
                    )
                    if not save_ok:
                        logger.warning(f"用户[{account_no}]的对话保存失败")
                    else:
                        logger.info(f"用户[{account_no}]的对话已保存（响应长度：{len(response)}字）")

                    # 成功处理后跳出重试循环
                    retry_count = max_retries
                    break

        # 检查最终结果
        if retry_count >= max_retries and not has_sent_data:
            error_msg = f"用户[{account_no}]的智能体连续{max_retries}次返回空响应，无法生成有效回答"
            logger.error(error_msg)
            raise ApiException(msg="抱歉，未能生成有效回答，请稍后再试")

        asyncio.create_task(chat_service.generate_summary(account_no=account_no))
        logger.debug(f"用户[{account_no}]的摘要更新任务已提交至后台")

    except Exception as e:
        error_msg = f"用户[{account_no}]的智能体对话失败：{str(e)}"
        logger.error(error_msg, exc_info=True)

        if has_sent_data:
            error_end_token = JsonData.stream_data(
                data="",
                msg="error"
            ).model_dump_json()
            yield f"data:{error_end_token}\n\n"
            yield "data:{\"code\":1,\"data\":\"\",\"msg\":\"[ERROR_DONE]\"}\n\n"
            await asyncio.sleep(0.01)
            raise StopAsyncIteration()
        else:
            raise ApiException(msg="智能体对话失败") from e


async def generate_stream_response(
        chat_service,
        message: str,
        account_no: str
) -> AsyncIterator[str]:
    setup_tracer()

    processed_message = apply_text_replacements(message)
    logger.info(f"用户[{account_no}]原始输入：{message} → 处理后：{processed_message}")

    knowledge_tool_list = await get_knowledge_tools()
    chat_tool_list = get_chat_tools()

    if not isinstance(knowledge_tool_list, list):
        raise TypeError(f"knowledge_tool_list 应为列表，实际是 {type(knowledge_tool_list)}")
    if not isinstance(chat_tool_list, list):
        raise TypeError(f"chat_tool_list 应为列表，实际是 {type(chat_tool_list)}")

    all_tools = knowledge_tool_list + chat_tool_list
    logger.info(
        f"已加载工具（共{len(all_tools)}个）："
        f"{[getattr(t, 'name', f'{t.__class__.__name__}（无name属性）') for t in all_tools]}"
    )

    agent = create_small_chat_agent(tools=all_tools)

    current_chunk = ""
    filter_symbols = ["*", "-", "=", "~", "`"]
    filter_combinations = ["**", "--", "=="]

    async for token in small_chat_with_agent(
            agent_executor=agent,
            account_no=account_no,
            chat_service=chat_service,
            processed_input=processed_message
    ):
        current_chunk += token
        if token in ["。", "？", "！", ";", "，", "、"] or len(current_chunk) >= 15:
            stream_data = JsonData.stream_data(
                data=current_chunk,
                msg="streaming"
            )
            yield f"data:{stream_data.model_dump_json()}\n\n"
            current_chunk = ""
            await asyncio.sleep(0.01)

    if current_chunk:
        for combo in filter_combinations:
            current_chunk = current_chunk.replace(combo, "")
        stream_data = JsonData.stream_data(
            data=current_chunk,
            msg="streaming"
        )
        yield f"data:{stream_data.model_dump_json()}\n\n"

    yield "data:{\"code\":0,\"data\":\"\",\"msg\":\"[DONE]\"}\n\n"
    logger.info(f"用户[{account_no}]的流式响应已完成")